/*
Copyright (C) Alex S. (mailto: md5sum@alexsp.ru)

This program is free software; you can redistribute it and/or
modify it under the terms of the GNU General Public License
as published by the Free Software Foundation; either version 2
of the License.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details
http://www.gnu.org/licenses/gpl-2.0.html.

You should have received a copy of the GNU General Public License
along with this program in the file gpl-2.0.txt;
if not, write to the Free Software Foundation,
Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
 */

#include <avr/io.h>
#include <avr/boot.h>
#include <avr/pgmspace.h>
#include <avr/eeprom.h>
#include "stk500protocol.h"
#include "glcd_boot.h"

uint16_t used_load_addr=0;
uint16_t used_read_addr=0;

#define BUFF_SIZE 285
static uchar buff[BUFF_SIZE];
uchar stk_communicate=1;

static uchar stk_params[]={
		0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
		STK_HW_VER,STK_SW_MAJOR,STK_SW_MINOR, /*0x90,0x91,0x92*/
		0x0,0x32,0x32,0x1,0x0,
		0x0,0x0,0xaa,0x0,0x0,0x0,0x1,0x1
};

#define STK_PAGE_MASK (0xFFFF&(SPM_PAGESIZE-1))

static uint16_t addr=0;
static const uchar sign[]={SIGNATURE_0,SIGNATURE_1,SIGNATURE_2};

#define _MESS ((stk_mess_t *)buff)
#define _MESS_BODY ((stk_mess_t *)buff)->body
#define _MESS_CMD ((stk_mess_t *)buff)->cmd

static void send_resp(const uint len) {
	uint i;
	uchar chksum=0;

	_MESS->sizeH=(uint8_t)((len>>8)&0xff);
	_MESS->sizeL=(uint8_t)(len&0xff);
	for (i=0;i<(len+5);i++) {
		chksum^=buff[i];
	}
	buff[len+5]=chksum;
	for (i=0;i<(len+6);i++) {
		while ((UCSRA&(1<<UDRE))==0) {};
		UDR=buff[i];
	}
}

static uint do_CMD_PROGRAM_FLASH_ISP(){
	uint i=0;
	uint len;

	len=(_MESS_BODY[0]<<8)+_MESS_BODY[1];

	if ((_MESS_BODY[2]&0x01)==0) {_MESS_BODY[0]=STATUS_CMD_FAILED;return 2;}

	do{
		if ((addr&STK_PAGE_MASK)==0) {
		    boot_spm_busy_wait();
		    eeprom_busy_wait();
			boot_page_erase((addr)&(~(STK_PAGE_MASK)));
		    boot_spm_busy_wait();
		    eeprom_busy_wait();
		}

		boot_page_fill((addr&STK_PAGE_MASK),((_MESS_BODY[i+10]<<8)|_MESS_BODY[i+9]));
		addr+=2;
		if ((addr>>10)+1>used_load_addr) {
			used_load_addr=(addr>>10)+1;
			GLCD_xy(load_x,7);
			print_load_addr(used_load_addr);
		}
		i+=2;
		if ((addr&STK_PAGE_MASK)==0) {
				boot_page_write_safe((addr-1)&(~(STK_PAGE_MASK)));
				boot_rww_enable_safe();
			}
	} while (i<len);
	if ((addr&STK_PAGE_MASK)!=0) if (_MESS_BODY[2]&0x80) {
			boot_page_write_safe((addr-1)&(~(STK_PAGE_MASK)));
			boot_rww_enable_safe();
		}
	_MESS_BODY[0]=STATUS_CMD_OK;
	return 2;
}
static uint do_CMD_READ_FLASH_ISP() {
	uint i,len;
	len=(_MESS_BODY[0]<<8)+_MESS_BODY[1];
	_MESS_BODY[0]=STATUS_CMD_OK;
	for (i=0;i<len;i++) {
		_MESS_BODY[i+1]=pgm_read_byte(addr++);
		if ((addr>>10)+1>used_read_addr) {
			used_read_addr=(addr>>10)+1;
			GLCD_xy(read_x,7);
			print_read_addr(used_read_addr);
		}
	}
	_MESS_BODY[len+1]=STATUS_CMD_OK;
	return len+3;
}
static uint do_CMD_PROGRAM_EEPROM_ISP(){
	uint len;
	len=(_MESS_BODY[0]<<8)+_MESS_BODY[1];
	eeprom_write_block((void *)(&_MESS_BODY[9]),(void *)addr,len);
	addr+=len;
	_MESS_BODY[0]=STATUS_CMD_OK;
	return 2;
}
static uint do_CMD_READ_EEPROM_ISP() {
	uint len;
	len=(_MESS_BODY[0]<<8)+_MESS_BODY[1];
	_MESS_BODY[0]=STATUS_CMD_OK;
	eeprom_read_block((void *)(&_MESS_BODY[1]),(void *)addr,len);
	addr+=len;
	_MESS_BODY[len+1]=STATUS_CMD_OK;
	return len+3;
}
static uchar get_fuse(uchar * p) {
	uchar ret=0;

	if ((*p==0x50)&&(*(p+1)==0x00))
			ret=boot_lock_fuse_bits_get(GET_LOW_FUSE_BITS);

	else if ((*p==0x58)&&(*(p+1)==0x08))
			ret=boot_lock_fuse_bits_get(GET_HIGH_FUSE_BITS);

	else if ((*p==0x58)&&(*(p+1)==0x00))
			ret=boot_lock_fuse_bits_get(GET_LOCK_BITS);
#if FUSE_MEMORY_SIZE>2
	else if ((*p==0x50)&&(*(p+1)==0x08))
			ret=boot_lock_fuse_bits_get(GET_EXTENDED_FUSE_BITS);
#endif
	return ret;
}
static uint do_CMD_SPI_MULTI() {
	uchar cmd1,cmd2,res=0;
	cmd1=_MESS_BODY[3];
	cmd2=_MESS_BODY[4];
	if (cmd1==0x30) {
		res=sign[(_MESS_BODY[5]&0x03)];
	}

	else if (cmd1==0xa0) {
		res=eeprom_read_byte((uint8_t *)(((_MESS_BODY[4]&3)<<8)|_MESS_BODY[5]));
	}

	else if (cmd1==0xc0) {
		eeprom_write_byte((uint8_t *)(((_MESS_BODY[4]&3)<<8)|_MESS_BODY[5]),_MESS_BODY[6]);
	}

	else if ((cmd1==0x50)||(cmd1==0x58)) res=get_fuse(&_MESS_BODY[3]);

	_MESS_BODY[0]=STATUS_CMD_OK;
	_MESS_BODY[1]=0;
	_MESS_BODY[2]=cmd1;
	_MESS_BODY[3]=cmd2;
	_MESS_BODY[4]=res;
	_MESS_BODY[5]=STATUS_CMD_OK;
	return 7;
}

uchar stk500_go(){
	uint resp=2;
	if (_MESS_CMD==CMD_SIGN_ON) {
//		stk_communicate=1;
		_MESS_BODY[0]=STATUS_CMD_OK;
		_MESS_BODY[1]=8;
		_MESS_BODY[2]='S';
		_MESS_BODY[3]='T';
		_MESS_BODY[4]='K';
		_MESS_BODY[5]='5';
		_MESS_BODY[6]='0';
		_MESS_BODY[7]='0';
		_MESS_BODY[8]='_';
		_MESS_BODY[9]='2';
		resp=11;
	}

	else if (_MESS_CMD==CMD_SET_PARAMETER) {
		stk_params[_MESS_BODY[0]&0x1f]=_MESS_BODY[1];
		_MESS_BODY[0]=STATUS_CMD_OK;
	}

	else if (_MESS_CMD==CMD_GET_PARAMETER) {
		_MESS_BODY[1]=stk_params[_MESS_BODY[0]&0x1f];
		_MESS_BODY[0]=STATUS_CMD_OK;
		resp=3;
	}

	else if (_MESS_CMD==CMD_OSCCAL) {_MESS_BODY[0]=STATUS_CMD_OK;/*resp=2;*/}

	else if (_MESS_CMD==CMD_LOAD_ADDRESS) {
		addr=(_MESS_BODY[2]<<8)+_MESS_BODY[3];
		addr<<=1;
		_MESS_BODY[0]=STATUS_CMD_OK;
	}

	else if (_MESS_CMD==CMD_FIRMWARE_UPGRADE) {_MESS_BODY[0]=STATUS_CMD_FAILED;/*resp=2;*/}

	else if (_MESS_CMD==CMD_ENTER_PROGMODE_ISP) {
//		stk_communicate=1;
		_MESS_BODY[0]=STATUS_CMD_OK;
	}

	else if (_MESS_CMD==CMD_LEAVE_PROGMODE_ISP) {
		stk_communicate=0;
		_MESS_BODY[0]=STATUS_CMD_OK;
	}

	else if (_MESS_CMD==CMD_CHIP_ERASE_ISP) {_MESS_BODY[0]=STATUS_CMD_OK;/*resp=2;*/}

	else if (_MESS_CMD==CMD_PROGRAM_FLASH_ISP) resp=do_CMD_PROGRAM_FLASH_ISP();

	else if (_MESS_CMD==CMD_READ_FLASH_ISP) resp=do_CMD_READ_FLASH_ISP();

	else if (_MESS_CMD==CMD_PROGRAM_EEPROM_ISP) resp=do_CMD_PROGRAM_EEPROM_ISP();

	else if (_MESS_CMD==CMD_READ_EEPROM_ISP) resp=do_CMD_READ_EEPROM_ISP();

	else if (_MESS_CMD==CMD_PROGRAM_FUSE_ISP) {_MESS_BODY[0]=STATUS_CMD_OK;_MESS_BODY[1]=STATUS_CMD_OK;resp=3;}

	else if ((_MESS_CMD==CMD_READ_FUSE_ISP)||(_MESS_CMD==CMD_READ_LOCK_ISP)){
/*
		if ((_MESS_BODY[1]==0x50)&&(_MESS_BODY[2]==0x00))
				_MESS_BODY[1]=boot_lock_fuse_bits_get(GET_LOW_FUSE_BITS);

		else if ((_MESS_BODY[1]==0x58)&&(_MESS_BODY[2]==0x08))
				_MESS_BODY[1]=boot_lock_fuse_bits_get(GET_HIGH_FUSE_BITS);

		else if ((_MESS_BODY[1]==0x58)&&(_MESS_BODY[2]==0x00))
				_MESS_BODY[1]=boot_lock_fuse_bits_get(GET_LOCK_BITS);
#if FUSE_MEMORY_SIZE>2
		else if ((_MESS_BODY[1]==0x50)&&(_MESS_BODY[2]==0x08))
				_MESS_BODY[1]=boot_lock_fuse_bits_get(GET_EXTENDED_FUSE_BITS);
#endif
*/
		_MESS_BODY[1]=get_fuse(&_MESS_BODY[1]);
		_MESS_BODY[0]=STATUS_CMD_OK;
		_MESS_BODY[2]=STATUS_CMD_OK;
		resp=4;
	}

	else if (_MESS_CMD==CMD_PROGRAM_LOCK_ISP) {
//		boot_lock_bits_set(_MESS_BODY[3]);
		_MESS_BODY[0]=STATUS_CMD_OK;
		_MESS_BODY[1]=STATUS_CMD_OK;
		resp=3;
	}

	else if (_MESS_CMD==CMD_READ_SIGNATURE_ISP) {
		_MESS_BODY[1]=sign[(_MESS_BODY[3]&0x03)];
		_MESS_BODY[0]=STATUS_CMD_OK;
		_MESS_BODY[2]=STATUS_CMD_OK;
		resp=4;
	}

	else if (_MESS_CMD==CMD_READ_OSCCAL_ISP) {
		_MESS_BODY[1]=0;
		_MESS_BODY[0]=STATUS_CMD_OK;
		_MESS_BODY[2]=STATUS_CMD_OK;
		resp=4;
	}

	else if (_MESS_CMD==CMD_SPI_MULTI) resp=do_CMD_SPI_MULTI();

	else {_MESS_BODY[0]=STATUS_CMD_UNKNOWN;/*resp=2;*/}

	send_resp(resp);
	return resp;
}
void send_chk_err(){
	_MESS->cmd=ANSWER_CKSUM_ERROR;
	_MESS_BODY[0]=STATUS_CKSUM_ERROR;
	send_resp(2);
}

void stk_buff_add(uchar b) {
	static uint ind=0;
	static uint len=6;
	static uchar chksum=0;

	if (ind==0) {
		if (b==MESSAGE_START) {
			buff[ind++]=b;
			chksum=b;
			len=6;
		}
		return;
	}
	if (ind<len-1) {
		buff[ind++]=b;
		chksum^=b;
	}else {
		if (chksum==b) stk500_go(); else send_chk_err();
		ind=0;
	}
	if (ind==5) {
		if (_MESS->token!=TOKEN) {
			send_chk_err();
			ind=0;
		}
		len=((_MESS->sizeH)<<8)+_MESS->sizeL+6;
	}
}
